#!/usr/bin/env python

"""
zbreplay: replay ZigBee/802.15.4 network traffic from libpcap or Daintree files
jwright@willhackforsushi.com
"""

import sys
import time
import signal
import argparse

from killerbee import *

def interrupt(signum, frame):
    global packetcount
    global kb
    global cap
    kb.close()
    if cap:
        cap.close()
    print("{0} packets transmitted".format(packetcount))
    sys.exit(0)

def packet_ack(packet):
    '''Returns True if the packet is an ACK.'''
    MAC_FC_FTYPE_MASK = 0x0007
    MAC_FC_FTYPE_ACK = 2
    fc = struct.unpack("H", packet[0:2])[0]
    if (fc & MAC_FC_FTYPE_MASK) == MAC_FC_FTYPE_ACK:
        return True
    return False

# Pcap or Daintree reader object
cap = None

# Global
packetcount = 0

# Command-line arguments
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument('-i', '--iface', '--dev', action='store', dest='devstring')
parser.add_argument('-r', '--pcapfile', action='store', default=None)
parser.add_argument('-R', '--dsnafile', action='store', default=None)
#parser.add_argument('-g', '--gps', '--ignore', action='append', dest='ignore')
parser.add_argument('-c', '-f', '--channel', action='store', type=int, default=None)
parser.add_argument('-n', '--count', action='store', type=int, default=-1)
parser.add_argument('-s', '--sleep', action='store', type=float, default=1.0)
parser.add_argument('-D', action='store_true', dest='showdev')
args = parser.parse_args()

if args.showdev:
    show_dev()
    sys.exit(0)

if args.channel == None:
    print >>sys.stderr, "ERROR: Must specify a channel."
    sys.exit(1)

if args.pcapfile == None and args.dsnafile == None:
    print >>sys.stderr, "ERROR: Must specify a file with frames to replay using -r (libpcap) or -R (Daintree SNA)"
    sys.exit(1)
elif args.pcapfile is not None:
    savefile = args.pcapfile
    cap = PcapReader(args.pcapfile)
elif args.dsnafile is not None:
    savefile = args.dsnafile
    cap = DainTreeReader(args.dsnafile)

kb = KillerBee(device=args.devstring)
signal.signal(signal.SIGINT, interrupt)
if not kb.is_valid_channel(args.channel):
    print >>sys.stderr, "ERROR: Must specify a valid IEEE 802.15.4 channel for the selected device."
    kb.close()
    sys.exit(1)
kb.set_channel(args.channel)

print("zbreplay: retransmitting frames from \'{0}\' on interface \'{1}\' with a delay of {2} seconds.".format(savefile, kb.get_dev_info()[0], args.sleep))

while args.count != packetcount:
    try:
        packet = cap.pnext()[1]
        # We don't want to replay ACK packets from the capture, typically.
        if not packet_ack(packet):
            packetcount += 1
            kb.inject(packet)
            time.sleep(args.sleep)
    except TypeError:       # raised when pnext returns Null (end of capture)
        break

kb.close()
cap.close()

print("{0} packets transmitted".format(packetcount))
